import torch
# tril = torch.tril(torch.ones(3,3))
# print(tril)
# print(1-tril)
def get_padding_mask(x,padding_idx):
    return (x == padding_idx).unsqueeze(1)

def get_subsequent_mask(size):
    mask_shape = (1,size,size)
    return 1-torch.tril(torch.ones(mask_shape)).byte()

inputs = torch.tensor([
    [1,2,3],
    [4,5,0]
])

pad_mask = get_padding_mask(inputs,0)
print(pad_mask)
sub_mask = get_subsequent_mask(3)
print(sub_mask)
mask = sub_mask | pad_mask
print(mask)